import os
import subprocess
import argparse
import os
import argparse
from simnibs import sim_struct, run_simnibs

def create_head_model(T1_path, T2_path=None, dMRI_path=None, subject_name=None, 
                      head_models_folder=None):
    """
    Creates a head model for use in SimNIBS, based on provided MRI and dMRI images.
    
    Parameters:
    - T1_path: str, required
        Path to the T1-weighted MRI file.
    - T2_path: str, optional
        Path to the T2-weighted MRI file.
    - dMRI_path: str, optional
        Path to the diffusion MRI (dMRI) file.
    - subject_name: str, optional
        Name of the subject. If not provided, it will be derived from the T1 file name.
    - head_models_folder: str, optional
        Folder where the head model will be saved. Defaults to the current working directory.
        
    Outputs:
    - m2m_{subject_name} folder is created in the specified head_models_folder.
    """
    
    # Use the current working directory if head_models_folder is not provided
    if head_models_folder is None:
        head_models_folder = os.getcwd()
    
    print(f"head_models_folder: {head_models_folder}")
    print(f"Will save outputs at: {os.path.join(head_models_folder, 'm2m_' + subject_name)}")
    
    # If subject_name is not provided, derive it from the T1_path filename
    if subject_name is None:
        subject_name = os.path.basename(T1_path).split(".")[0]
    
    # Check if the subject directory already exists
    if os.path.exists(os.path.join(head_models_folder, "m2m_" + subject_name)):
        print(f"Subject '{subject_name}' already exists, try another name or assign a different name")
        return 0
    #===========================================================================================================#
    os.chdir(head_models_folder)

    print(f"Subject name: {subject_name}")
    
    # Construct the command for SimNIBS charm
    commands = f"charm {subject_name} {T1_path}"
    print(f"Initial command: {commands}")
    
    # If a T2 file is provided, append it to the command
    if T2_path is not None:
        commands += f" {T2_path}"
    print(f"Updated command with T2 (if provided): {commands}")

    # Execute the command for the head model creation
    try:
        subprocess.run(commands, check=True, shell=True)
        print("Head model creation command executed successfully.")
    except subprocess.CalledProcessError as e:
        
        raw_commands=commands
        print(f"Error executing head model creation command: {e}")
        try:
            commands=raw_commands+" --forcerun"
            subprocess.run(commands, check=True, shell=True)
            print("Head model creation command executed successfully.")
        except subprocess.CalledProcessError as e:
            print(f"Error executing head model creation command: {e}")
            try:
                #commands += " --forceqform --forcerun" 
                commands=raw_commands+" --forceqform --forcerun"
                subprocess.run(commands, check=True, shell=True)
                print("Head model creation command executed successfully.")
            except subprocess.CalledProcessError as e:
                print(f"Error executing head model creation command: {e}")
                try:
                    #commands += " --forcesform --forcerun"
                    commands=raw_commands+" --forcesform --forcerun"
                    subprocess.run(commands, check=True, shell=True)
                    print("Head model creation command executed successfully.")
                except subprocess.CalledProcessError as e:
                    print(f"Error executing head model creation command: {e}")
    #===========================================================================================================#
    # Process dMRI if provided
    if dMRI_path is not None:
        # Determine paths for bval, bvec, and reverse-phase-encoded images
        bval_path = dMRI_path.replace(".nii.gz", ".bval")
        bvec_path = dMRI_path.replace(".nii.gz", ".bvec")
        rev_path = dMRI_path.replace(".nii.gz", "_rev.nii.gz")
        
        # Build the second command for dMRI processing
        #command2 = f"dwi2cond --all {subject_name} {dMRI_path} {bval_path} {bvec_path} {rev_path}"
        command2 = f"dwi2cond --all {subject_name} {dMRI_path}"
        if os.path.exists(bval_path):
            command2 += f" {bval_path}"
        if os.path.exists(bvec_path):
            command2 += f" {bvec_path}"
        if os.path.exists(rev_path):
            command2 += f" {rev_path}"
        print(f"dMRI processing command: {command2}")
        
        # Execute the command for dMRI processing
        try:
            subprocess.run(command2, check=True, shell=True)
            print("dMRI processing command executed successfully.")
        except subprocess.CalledProcessError as e:
            print(f"Error executing dMRI processing command: {e}")
    
        # Return to the original directory after processing
        print(f"Returning to original directory: {os.getcwd()}")
    print("Creating head model finished.")
    #===========================================================================================================#
    print("*"*50)
    print("Creating leadfield...")
    from tdcs_leadfield import run_leadfield_simulation
    run_leadfield_simulation("m2m_"+subject_name, leadfield_name=None)
    print("Creating leadfield finished")
    #===========================================================================================================#
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Create a head model using SimNIBS based on MRI and dMRI images.")
    
    # Adding arguments
    parser.add_argument("-T1_path", required=True, help="Path to the T1-weighted MRI file.")
    parser.add_argument("-T2_path", help="Path to the T2-weighted MRI file (optional).")
    parser.add_argument("-dMRI_path", help="Path to the diffusion MRI (dMRI) file (optional).")
    parser.add_argument("-subject_name", help="Name of the subject (optional).")
    parser.add_argument("-head_models_folder", help="Folder where the head model will be saved (optional).")
    
    # Parsing arguments
    args = parser.parse_args()
    
    # Calling the function with parsed arguments
    create_head_model(
        T1_path=args.T1_path,
        T2_path=args.T2_path,
        dMRI_path=args.dMRI_path,
        subject_name=args.subject_name,
        head_models_folder=args.head_models_folder
    )
